package com.da.llm;

import com.da.message.Message;
import com.da.prompts.PromptTemplate;
import com.da.result.*;
import com.da.tools.Parse;
import com.da.tools.Request;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Ollama的调用类
 *
 * @author da
 * @time 2024/7/3 上午 11:12
 */
public class Ollama {
    /**
     * Ollama的请求地址
     */
    private final String url;
    /**
     * 要使用的模型名称
     */
    private final String model;
    /**
     * 模型的温度,范围从0到1,用于控制生成文本的随机性
     */
    private final float temperature;
    /**
     * 聊天历史列表
     */
    private final List<Message> history = new ArrayList<>();

    /**
     * 构造函数
     *
     * @param url         请求地址
     * @param model       模型名称
     * @param temperature 温度
     */
    public Ollama(String url, String model, float temperature) {
        this.url = url;
        this.model = model;
        this.temperature = temperature;
    }

    /**
     * 构造函数
     * 默认请求地址为 localhost:11434
     * 温度为 0.8f
     *
     * @param model 模型名称
     */
    public Ollama(String model) {
        this("http://localhost:11434", model, 0.8f);
    }

    /**
     * 构造函数
     * 默认请求地址为 localhost:11434
     *
     * @param model       模型名称
     * @param temperature 温度
     */
    public Ollama(String model, float temperature) {
        this("http://localhost:11434", model, temperature);
    }

    /**
     * 构造函数
     * 默认请求地址为 localhost:11434
     * 默认模型为 deepseek-r1:1.5b
     * 温度为 0.8f
     */
    public Ollama() {
        this("http://localhost:11434", "deepseek-r1:1.5b", 0.8f);
    }

    /**
     * 调用生成回答接口
     *
     * @param prompt 问题
     * @return 模型响应结果
     */
    public Result run(String prompt) {
        return generate(prompt);
    }

    /**
     * 调用生成回答接口
     *
     * @param prompt 问题prompt
     * @return 模型响应结果
     */
    public Result run(PromptTemplate prompt) {
        return run(prompt.toString());
    }

    /**
     * 调用生成回答接口
     *
     * @param prompt 问题prompt
     * @param args   模板对应的键值
     * @return 模型响应结果
     */
    public Result run(PromptTemplate prompt, Object... args) {
        return run(prompt.format(args));
    }

    /**
     * 调用聊天回答接口
     *
     * @param history 聊天历史
     * @return 模型响应结果
     */
    public Result chat(List<Message> history) {
        this.history.addAll(history);
        Request request = createRequest("chat", this.history);
        return handleRequest("chat", request);
    }

    /**
     * 调用聊天回答接口
     *
     * @param message 角色+问题
     * @return 模型响应结果
     */
    public Result chat(Message message) {
        //保存当前的问题到历史记录
        this.history.add(message);
        Request request = createRequest("chat", this.history);
        return handleRequest("chat", request);
    }

    /**
     * 请求接口生成回答
     *
     * @param prompt 问题
     * @return 模型响应结果
     */
    private Result generate(String prompt) {
        Request request = createRequest("generate", prompt);
        return handleRequest("generate", request);
    }

    /**
     * 处理响应请求
     *
     * @param type    区分生成回答和聊天回答
     * @param request 请求对象
     * @return 模型响应结果
     */
    private Result handleRequest(String type, Request request) {
        String response = request.send();
        if (request.getCode() == 200) {
            switch (type) {
                case "generate":
                    return Parse.fromGenerateJson(response);
                case "chat":
                    ChatResult chatResult = Parse.fromChatJson(response);
                    //保存当前的响应到历史记录
                    this.history.add(chatResult.getMessage());
                    return chatResult;
                case "tags":
                    return Parse.fromTagsJson(response);
                default:
                    break;
            }
        }
        return new EmptyResult(request.getCode());
    }

    /**
     * 创建请求对象
     *
     * @param type   区分生成回答和聊天回答
     * @param prompt 问题或者历史消息列表
     * @return 请求对象
     */
    private Request createRequest(String type, Object prompt) {
        String url = this.url + "/api/" + type;
        Map<String, Object> data = new HashMap<>();
        data.put("model", this.model);
        data.put("temperature", this.temperature);
        switch (type) {
            case "generate":
                data.put("stream", false);
                data.put("prompt", prompt);
                break;
            case "generate_stream":
                //去掉请求地址的_stream后缀
                url = url.replace("_stream", "");
                data.put("stream", true);
                data.put("prompt", prompt);
                break;
            case "chat":
                data.put("stream", false);
                data.put("messages", prompt);
                break;
            case "chat_stream":
                //去掉请求地址的_stream后缀
                url = url.replace("_stream", "");
                data.put("stream", true);
                data.put("messages", prompt);
                break;
            default:
                break;
        }
        return new Request(url, data);
    }

    /**
     * 获取当前会话的历史消息列表
     *
     * @return 当前会话保存的历史消息列表
     */
    public List<Message> chat_history() {
        return this.history;
    }

    /**
     * 获取ollama安装到本地的模型信息
     *
     * @return 模型信息列表
     */
    public ModelInfoResult list() {
        Request request = new Request(this.url + "/api/tags", "GET", null);
        return (ModelInfoResult) handleRequest("tags", request);
    }

    /**
     * 获取模型信息
     *
     * @param name 模型名字
     * @return 模型信息
     */
    public String show(String name) {
        Map<String, Object> data = new HashMap<>();
        data.put("name", name);
        Request request = new Request(this.url + "/api/show", data);
        String response = request.send();
        return request.getCode() == 200 ? response : "";
    }

    /**
     * 调用流式生成回答接口
     *
     * @param prompt 问题prompt
     * @return 模型响应结果
     */
    public StreamResult run_stream(String prompt) {
        //指定生成流式响应
        Request request = createRequest("generate_stream", prompt);
        return request.stream("generate");
    }

    /**
     * 调用生成回答接口(流式生成)
     *
     * @param prompt 问题prompt
     * @return 模型响应结果
     */
    public StreamResult run_stream(PromptTemplate prompt) {
        return run_stream(prompt.toString());
    }

    /**
     * 调用生成回答接口(流式生成)
     *
     * @param prompt 问题prompt
     * @param args   模板对应的键值
     * @return 模型响应结果
     */
    public StreamResult run_stream(PromptTemplate prompt, Object... args) {
        return run_stream(prompt.format(args));
    }

    /**
     * 流式调用聊天回答接口
     *
     * @param message 角色+问题
     * @return 模型响应结果
     */
    public StreamResult chat_stream(Message message) {
        //保存当前的问题到历史记录
        this.history.add(message);
        Request request = createRequest("chat_stream", this.history);
        return request.stream("chat");
    }

    /**
     * 调用聊天回答接口
     *
     * @param history 聊天历史
     * @return 模型响应结果
     */
    public StreamResult chat_stream(List<Message> history) {
        this.history.addAll(history);
        Request request = createRequest("chat_stream", this.history);
        return request.stream("chat");
    }
}
